import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression
import pandas as pd
import unicodedata
file_path = r"C:\Users\Lenovo User\Desktop\美赛\2025_Problem_C_Data\summerOly_medal_counts.csv"  # 确保文件是CSV格式

# 读取文件，只提取 NOC 列
df = pd.read_csv(file_path, usecols=["NOC"])


# 去除特殊字符
def normalize(text):
    return unicodedata.normalize("NFKC", text).strip()
countries = [normalize(country) for country in df["NOC"].unique()]
to_remove = [
    'Soviet Union', 'East Germany', 'West Germany', 'Unified Team', 'Yugoslavia',
    'Czechoslovakia', 'Macedonia', 'ROC', 'FR Yugoslavia','Mixed team',
    'Australasia', 'Bohemia', 'British West Indies', 'Independent Olympic Participants',
    'Russian Empire','Formosa', 'British West Indies', 'United Team of Germany', 'Soviet Union',
      'Taiwan', 'Chinese Taipei', 'Netherlands Antilles', 'Virgin Islands', 'Hong Kong', 
      'Independent Olympic Athletes', 'Refugee Olympic Team', 'Kosovo',
      'Bermuda', 
      'Puerto Rico','Ceylon'
]
to_remove = [normalize(country) for country in to_remove]
countries = [country for country in countries if country not in to_remove]


# 打印结果
print("Unique countries:", countries)

df = pd.read_csv(file_path)

# 确保列名一致
df.columns = ["Rank", "Country", "Gold", "Silver", "Bronze", "Total", "Year"]
pd.set_option('display.max_rows',None)
print("Sample Data:\n", df)
import numpy as np
import pandas as pd
from sklearn.preprocessing import MinMaxScaler
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense, Dropout
from sklearn.metrics import mean_squared_error, r2_score
import tensorflow as tf



def forecast_medals_lstm(country, df, end_year=2028, look_back=3):
    """使用LSTM预测奖牌数并计算评估指标"""
    try:
        # 1. 数据提取和验证
        country_data = df[df["Country"] == country]
        if len(country_data) < look_back + 1:
            print(f"警告: {country} 数据量不足")
            return None
            
        # 2. 数据准备
        gold_data = country_data["Gold"].values.reshape(-1, 1)
        total_data = country_data["Total"].values.reshape(-1, 1)
        
        # 3. 数据验证
        for data, name in [(gold_data, "金牌"), (total_data, "总奖牌")]:
            if np.isnan(data).any() or np.isinf(data).any():
                print(f"警告: {country} {name}数据异常")
                return None
        
        # 4. 数据标准化
        scaler_gold = MinMaxScaler()
        scaler_total = MinMaxScaler()
        
        gold_scaled = scaler_gold.fit_transform(gold_data)
        total_scaled = scaler_total.fit_transform(total_data)
        
        # 5. 准备训练数据
        def prepare_sequence(data):
            X, y = [], []
            for i in range(len(data) - look_back):
                X.append(data[i:(i + look_back)])
                y.append(data[i + look_back])
            return np.array(X), np.array(y)
            
        X_gold, y_gold = prepare_sequence(gold_scaled)
        X_total, y_total = prepare_sequence(total_scaled)
        
        # 6. 创建模型
        def create_model():
            model = Sequential([
                LSTM(50, activation='relu', input_shape=(look_back, 1), return_sequences=True),
                Dropout(0.2),
                LSTM(50, activation='relu'),
                Dropout(0.2),
                Dense(1)
            ])
            model.compile(optimizer='adam', loss='mse', run_eagerly=True)
            return model
            
        # 7. 训练模型
        gold_model = create_model()
        total_model = create_model()
        
        gold_model.fit(X_gold, y_gold, epochs=50, batch_size=1, verbose=0)
        total_model.fit(X_total, y_total, epochs=50, batch_size=1, verbose=0)
        
        # 8. 预测未来值
        future_years = np.arange(country_data["Year"].values[-1] + 4, end_year + 1, 4)
        
        def predict_sequence(model, last_sequence, scaler):
            predictions = []
            curr_seq = last_sequence.copy()
            
            for _ in range(len(future_years)):
                next_pred = model.predict(curr_seq.reshape(1, look_back, 1))
                predictions.append(next_pred[0])
                curr_seq = np.roll(curr_seq, -1)
                curr_seq[-1] = next_pred
            
            return scaler.inverse_transform(np.array(predictions).reshape(-1, 1))
            
        # 9. 生成预测
        last_gold = gold_scaled[-look_back:]
        last_total = total_scaled[-look_back:]
        
        gold_predictions = predict_sequence(gold_model, last_gold, scaler_gold)
        total_predictions = predict_sequence(total_model, last_total, scaler_total)
        
        # 10. 计算评估指标
        gold_pred_train = gold_model.predict(X_gold)
        total_pred_train = total_model.predict(X_total)
        
        # 计算MSE和R²
        gold_mse = mean_squared_error(y_gold, gold_pred_train)
        gold_r2 = r2_score(y_gold, gold_pred_train)
        total_mse = mean_squared_error(y_total, total_pred_train)
        total_r2 = r2_score(y_total, total_pred_train)
        
        # 11. 计算置信区间
        gold_std = np.std(gold_pred_train - y_gold)
        total_std = np.std(total_pred_train - y_total)
        
        z_score = 1.96  # 95% 置信区间
        gold_ci = (gold_predictions - z_score * gold_std, 
                  gold_predictions + z_score * gold_std)
        total_ci = (total_predictions - z_score * total_std,
                   total_predictions + z_score * total_std)
        
        # 12. 返回结果
        return {
            'future_years': future_years,
            'gold_pred': gold_predictions.flatten(),
            'total_pred': total_predictions.flatten(),
            'gold_ci': (gold_ci[0].flatten(), gold_ci[1].flatten()),
            'total_ci': (total_ci[0].flatten(), total_ci[1].flatten()),
            'metrics': {
                'gold_mse': gold_mse,
                'gold_r2': gold_r2,
                'total_mse': total_mse,
                'total_r2': total_r2
            }
        }
        
    except Exception as e:
        print(f"错误: 处理 {country} 时发生异常: {str(e)}")
        return None



# 使用示例
results_dict = {}
for country in countries:  
    results = forecast_medals_lstm(country, df)
    
    if results is not None:
        results_dict[country] = results
        print(f"\n{country} 预测结果:")
        print(f"预测年份: {results['future_years']}")
        
        print(f"\n金牌预测:")
        print(f"预测值: {results['gold_pred']}")
        print(f"95%置信区间: ({results['gold_ci'][0]}, {results['gold_ci'][1]})")
        print(f"MSE: {results['metrics']['gold_mse']:.4f}")
        print(f"R²: {results['metrics']['gold_r2']:.4f}")
        
        print(f"\n总奖牌预测:")
        print(f"预测值: {results['total_pred']}")
        print(f"95%置信区间: ({results['total_ci'][0]}, {results['total_ci'][1]})")
        print(f"MSE: {results['metrics']['total_mse']:.4f}")
        print(f"R²: {results['metrics']['total_r2']:.4f}")
    else:
        print(f"\n{country} 的预测被跳过。")

# 打印成功预测的国家数量
print(f"\n成功预测的国家数量: {len(results_dict)}")
import matplotlib.pyplot as plt
import matplotlib.font_manager as fm



def plot_predictions(countries, df, results_dict):
    """绘制预测结果图表"""
    plt.rcParams['font.sans-serif'] = ['SimHei'] # 使用黑体（SimHei），兼容 Windows
    plt.rcParams['axes.unicode_minus'] = False # 解决负号显示问题
    # 创建子图
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(15, 12))
    
    # 颜色循环
    colors = plt.cm.Set3(np.linspace(0, 1, len(countries)))
    
    # 绘制金牌预测
    for country, color in zip(countries, colors):
        if country in results_dict:
            results = results_dict[country]
            country_data = df[df["Country"] == country]
            
            # 历史数据
            ax1.plot(country_data["Year"], country_data["Gold"], 
                    color=color, label=f"{country} (历史)", marker='o')
            
            # 预测数据
            ax1.plot(results['future_years'], results['gold_pred'], 
                    color=color, linestyle='--', marker='x', markersize=10)
            
            # 置信区间
            ax1.fill_between(results['future_years'],
                           results['gold_ci'][0], results['gold_ci'][1],
                           color=color, alpha=0.2)
    
    ax1.set_xlabel("年份")
    ax1.set_ylabel("金牌数")
    ax1.set_title("奥运会金牌数预测 (2024-2028)")
    ax1.grid(True, alpha=0.3)
    ax1.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
    
    # 绘制总奖牌预测
    for country, color in zip(countries, colors):
        if country in results_dict:
            results = results_dict[country]
            country_data = df[df["Country"] == country]
            
            # 历史数据
            ax2.plot(country_data["Year"], country_data["Total"], 
                    color=color, label=f"{country} (历史)", marker='o')
            
            # 预测数据
            ax2.plot(results['future_years'], results['total_pred'], 
                    color=color, linestyle='--', marker='x', markersize=10)
            
            # 置信区间
            ax2.fill_between(results['future_years'],
                           results['total_ci'][0], results['total_ci'][1],
                           color=color, alpha=0.2)
    
    ax2.set_xlabel("年份")
    ax2.set_ylabel("总奖牌数")
    ax2.set_title("奥运会总奖牌数预测 (2024-2028)")
    ax2.grid(True, alpha=0.3)
    ax2.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
    
    plt.tight_layout()
    plt.show()

# 调用函数
plot_predictions(countries, df, results_dict)